import OpenAI from "openai";
import { getConfig } from "./config";

const getModelConfig = <T>(key: string) => getConfig<T>(`model.${key}`);

const chatTask = (client: OpenAI) => async (prompt: string) => {
    const completion = await client.chat.completions.create({
        model: getModelConfig("model_id"),
        messages: [
            { role: "system", content: "You are an expert at Rust programming." },
            { role: "user", content: prompt },
        ],
        stream: true,
        temperature: getModelConfig("temp"),
        top_p: getModelConfig("top_p"),
        stop: ["[/CODE]", "///"],
        max_tokens: getModelConfig("max_tokens"),
    });
    var ans = "";
    for await (const chunk of completion) {
        const content = chunk.choices[0].delta.content;
        if (content) {
            ans += content;
        }
    }
    return ans;
};

const cmplTask = (client: OpenAI) => async (prompt: string) => {
    const completion = await client.completions.create({
        model: getModelConfig("model_id"),
        prompt: prompt,
        echo: false,
        stream: true,
        temperature: getModelConfig("temp"),
        top_p: getModelConfig("top_p"),
        stop: ["[/CODE]", "///", "pub"],
        max_tokens: getModelConfig("max_tokens"),
    });
    var ans = "";
    for await (const chunk of completion) {
        const content = chunk.choices[0].text;
        if (content) {
            ans += content;
        }
    }
    return ans;
};

const invokeByModel = (client: OpenAI, modelId: string) =>
    modelId.startsWith("gpt") ? chatTask(client) : cmplTask(client);

const postProcess = (text: string) => {
    const start_idx = text.indexOf('[CODE]');
    if (start_idx !== -1) {
        text = text.substring(start_idx + 6).trim();
    }
    const end_idx = text.indexOf('[/CODE');
    if (end_idx !== -1) {
        text = text.substring(0, end_idx).trim();
    }
    return text;
};

export class PromptData {
    focal_ctx: string;
    rag_data: string;
    docstring: string;
    focal_fn_signature: string;

    constructor(v0: string, v1: string, v2: string, v3: string) {
        this.focal_ctx = v0;
        this.rag_data = v1;
        this.docstring = v2;
        this.focal_fn_signature = v3;
    }
}

export const buildPrompt = (data: PromptData) =>
    `You are given a blank Rust function with type context, referential code snippet and docstrings. Complete the Rust function with correct and functional code, and make use of the provided context if possible. Do NOT output any extra information. Provide your answer in [CODE] and [/CODE] tags.

[CONTEXT]
${data.focal_ctx}
${data.rag_data}
[/CONTEXT]
[CODE]
${data.docstring}
${data.focal_fn_signature}`;

export const invoke = async (prompt: string) => {
    const client = new OpenAI({
        baseURL: getModelConfig("base_url"),
        apiKey: getModelConfig("api_key"),
    });
    const modelId = getModelConfig<string>("model_id");
    return postProcess(await invokeByModel(client, modelId)(prompt));
};
